#include "waterfilling.h"

using namespace std;
using namespace cv;

cv::Mat WaterFilling::ApplyWaterFilling(cv::Mat & depthImage)
{
	Mat origin;
	depthImage.copyTo(origin);

	iterationNumber = 2;
 	waterInADrop = 4.0f;
 	waterDropped = 0.1f;
	scale = 5.f / (waterInADrop / waterDropped);
	factor = waterDropped;

	clock_t time;
	time = clock();

//	step = 0;
	GetMeasureFunction(depthImage);

	time = clock() - time;
	
	float tmp = float(time) / CLOCKS_PER_SEC;

	allTime += tmp;

	if (tmp > maxTime)
		maxTime = tmp;

	if (tmp < minTime)
		minTime = tmp;

	cout << "Time: " << tmp << endl;
	//cout << step << endl;
/*	getchar();*/

//  	for (int i = 0; i < depthImage.rows; ++i)
//  	{
//  		for (int j = 0; j < depthImage.cols; ++j)
//  		{
//  			auto mfit = measureFunction.find(make_pair(i, j));
//  			if (mfit == measureFunction.end())
//  			{
//  				depthImage.at<uchar>(i, j) = 0;
//  			}
//  			else
//  			{
//  				depthImage.at<uchar>(i, j) = mfit->second * 20;
//  			}
//  		}
//  	}
// 	measureFunction.clear();
//  	return depthImage;


	Thresholding(depthImage);

 	//measureFunction.clear();
 	//return depthImage;

	vector<vector<cv::Point>> contours = ContourAnalisys(depthImage);

// 	measureFunction.clear();
// 	return depthImage;

	if (contours.size() > 0)
	{
		cv::drawContours(origin, contours, -1, Scalar(255), CV_FILLED);
	}
	
	measureFunction.clear();
	return origin;
}

cv::Mat WaterFilling::ApplyFastWaterFilling(cv::Mat & depthImage)
{
	return depthImage;
}

float WaterFilling::ComputeMaximum()
{
	int histogram[256];
	float max = 0.0f;

	for (auto mfit = measureFunction.begin(); mfit != measureFunction.end(); ++mfit)
	{
		max = std::max(max, mfit->second);
	}
	cout << max << endl;
	for (int i = 0; i < 256; ++i)
	{
		histogram[i] = 0;
	}

	for (auto it = measureFunction.begin(); it != measureFunction.end(); ++it)
	{
		int index = it->second * 255.0f / max;
		histogram[min(index, 255)]++;
	}

	int noise = 0.02f * measureFunction.size();
	for (int i = 255; i >= 0; --i)
	{
		if (histogram[i] > noise)
		{
			max = float(i) / 255.0f * max;
			break;
		}
	}
	cout << max << endl;
	return max;
}


void WaterFilling::GetMeasureFunction(cv::Mat & depthImage)
{
	vector<pair<int, int>> foregroundDepth;

	//get indexes of foreground pixels in depthImage
	for (int i = 0; i < depthImage.rows; ++i)
	{
		for (int j = 0; j < depthImage.cols; ++j)
		{
			if (depthImage.at<uchar>(i, j) > 0)
				foregroundDepth.emplace_back(i, j);
		}
	}
	cout << "size" << foregroundDepth.size() << endl;

	//water filling algorithm
	//create distribution, but do we really need it?
	default_random_engine generator;
	uniform_int_distribution<int> indexDistribution(0, foregroundDepth.size() - 1);

	numberOfDrops = foregroundDepth.size();

	for (int k = 0; k < scale * numberOfDrops; ++k)
	{
		//get indexes of pixel to drop water in it
		//int x = foregroundDepth[k % numberOfDrops].first;
		//int y = foregroundDepth[k % numberOfDrops].second;
		int ind = indexDistribution(generator);
		int x = foregroundDepth[ind].first;
		int y = foregroundDepth[ind].second;
		int lx = 0, ly = 0;
		int iterations = iterationNumber;
		float remainWater = waterInADrop;
	//	factor = waterDropped;

		//int step_tmp = 0;

		while (remainWater > 0)
		{
			float min = 0.0f;
			map<pair<int, int>, float> neighbors;
			auto it = measureFunction.find(make_pair(x, y));
			float current = (it == measureFunction.end() ? 0 : it->second) + (depthImage.at<uchar>(x, y) == 0 ? 255 : depthImage.at<uchar>(x, y));
			//	float current = (it == measureFunction.end() ? 0 : it->second) + (255 - depthImage.at<uchar>(x, y));

		//	step_tmp++;

			//calculate values of neighbors
			for (int i = x - 1; i <= x + 1; ++i)
			{
				if (i < 0 || i >= depthImage.rows)
					continue;
				for (int j = y - 1; j <= y + 1; ++j)
				{
					if (j < 0 || j >= depthImage.cols)
						continue;
					auto nit = measureFunction.find(make_pair(i, j));
					float tmp = (nit == measureFunction.end() ? 0 : nit->second) + (depthImage.at<uchar>(i, j) == 0 ? 255 : depthImage.at<uchar>(i, j));
					//	float tmp = (nit == measureFunction.end() ? 0 : nit->second) + (255 - depthImage.at<uchar>(i, j));
					min = std::min(tmp - current, min);
					neighbors.emplace(make_pair(i, j), tmp - current);
				}
			}

			//find all pixels with min value
			float eps = 0.00001f;
			set<pair<int, int>> minValue;
			for (auto nit = neighbors.begin(); nit != neighbors.end(); ++nit)
			{
				if (fabs(nit->second - min) < eps)
					minValue.insert(nit->first);
			}


			if ((fabs(min + (waterDropped - factor) - 0.0f) < eps && (iterations == 0 || minValue.size() == 1))
				|| (min + (waterDropped - factor) > 0))
			{
				auto mit = measureFunction.find(make_pair(x, y));
				if (mit == measureFunction.end())
				{
					measureFunction.emplace(make_pair(x, y), std::min(waterDropped, remainWater));
				}
				else
				{
					measureFunction[make_pair(x, y)] += std::min(waterDropped, remainWater);
				}
				remainWater -= waterDropped;
				iterations = iterationNumber;
			//	factor = 0;
				continue;
			}

			if (fabs(min + (waterDropped - factor) - 0.0f) < eps && ((lx == 0 && ly == 0) || minValue.find(make_pair(x - lx, y - ly)) == minValue.end()))
			{
				uniform_int_distribution<int> minDist(0, minValue.size() - 1);
				int go = minDist(generator);
				auto vit = minValue.begin();
				while (go) { vit++; go--; }
				lx = x - vit->first;
				ly = y - vit->second;
				x = vit->first;
				y = vit->second;
				iterations--;
				continue;
			}

			//min value = 0 and there is a priority direction and we can move in this direction, 
			//so we move and priority direction is the same next step
			if (fabs(min + (waterDropped - factor) - 0.0f) < eps && minValue.find(make_pair(x - lx, y - ly)) != minValue.end())
			{
				x -= lx;
				y -= ly;
				iterations--;
				continue;
			}

			//min value != 0 and we have only one pixel or we don't have pixel in priority direction, we move to first pixel
			//and change priority direction
			if (min + (waterDropped - factor) < 0 && minValue.size() == 1)
			{
				auto vit = minValue.begin();
				lx = x - vit->first;
				ly = y - vit->second;
				x = vit->first;
				y = vit->second;
				iterations = iterationNumber;
				continue;
			}

			if (min + (waterDropped - factor) < 0 && minValue.find(make_pair(x - lx, y - ly)) == minValue.end())
			{
				uniform_int_distribution<int> minDist(0, minValue.size() - 1);
				int go = minDist(generator);
				auto vit = minValue.begin();
				while (go) { vit++; go--; }
				lx = x - vit->first;
				ly = y - vit->second;
				x = vit->first;
				y = vit->second;
				iterations = iterationNumber;
				continue;
			}

			//min value != 0 and we have pixel in priority direction, so we move to it and don't change priority direction
			if (min + (waterDropped - factor) < 0 && minValue.find(make_pair(x - lx, y - ly)) != minValue.end())
			{
				x -= lx;
				y -= ly;
				iterations = iterationNumber;
				continue;
			}
			cout << "BIG TROUBLE" << endl;
		}
		//if (step_tmp > step)
			//step = step_tmp;
	}
}


void WaterFilling::Thresholding(cv::Mat & depthImage)
{
	float max = 0.0f;

	for (auto mfit = measureFunction.begin(); mfit != measureFunction.end(); ++mfit)
	{
		max = std::max(max, mfit->second);
	}

	threshold = waterDropped * 10.0f * 255.0f / max;

	for (int i = 0; i < depthImage.rows; ++i)
	{
		for (int j = 0; j < depthImage.cols; ++j)
		{
			auto mfit = measureFunction.find(make_pair(i, j));
			if (mfit == measureFunction.end())
			{
				depthImage.at<uchar>(i, j) = 0;
			}
			else
			{

				float value = mfit->second * 255.0f / max;
				if (value >= threshold)
					depthImage.at<uchar>(i, j) = min(value, 255.0f);
				else
					depthImage.at<uchar>(i, j) = 0;
			}
		}
	}
}

// void WaterFilling::Thresholding(cv::Mat & depthImage)
// {
// 	float max = ComputeMaximum();
// 	float oneDrop = waterDropped * 255.0f / max;
// 	threshold = 255.0f - 6.0f * oneDrop;
// 
// 	for (int i = 0; i < depthImage.rows; ++i)
// 	{
// 		for (int j = 0; j < depthImage.cols; ++j)
// 		{
// 			auto mfit = measureFunction.find(make_pair(i, j));
// 			if (mfit == measureFunction.end())
// 			{
// 				depthImage.at<uchar>(i, j) = 0;
// 			}
// 			else
// 			{
// 
// 				float value = mfit->second * 255.0f / max;
// 				if (value >= threshold)
// 					depthImage.at<uchar>(i, j) = min(value, 255.0f);
// 				else
// 					depthImage.at<uchar>(i, j) = 0;
// 			}
// 		}
// 	}
// }

std::vector<std::vector<cv::Point>> WaterFilling::ContourAnalisys(cv::Mat & thresholdedImage)
{
	vector<vector<Point>> contours;

	Mat edges(cv::Size(thresholdedImage.rows, thresholdedImage.cols), CV_8U);
	cv::Canny(thresholdedImage, edges, 60, 180, 3);

	//edges.copyTo(thresholdedImage);
	//return contours;
	
	int type = MORPH_RECT;
	int size = 1;

	Mat element = getStructuringElement(type,
		Size(2 * size + 1, 2 * size + 1),
		Point(size, size));
	Mat closedEdges(cv::Size(thresholdedImage.rows, thresholdedImage.cols), CV_8U);
	cv::morphologyEx(edges, closedEdges, MORPH_CLOSE, element);
	//return closedEdges;

	
	cv::findContours(closedEdges, contours, CV_RETR_EXTERNAL, CV_CHAIN_APPROX_SIMPLE);
	//return contours;

	for (auto cit = contours.begin(); cit != contours.end();)
	{
		float area = cv::contourArea(*cit);
		float perimeter = cv::arcLength(*cit, true);
		float compact = area / std::pow(perimeter, 2);

		if (!(compact > 0.05 && compact < 0.08) || area < 100)
		{
			cit = contours.erase(cit);
			if (cit == contours.end())
				break;
			continue;
		}
		if (cit == contours.end())
			break;
		++cit;
	}
	return contours;
}

